#define IS_MAP_IMPLEMENTATION
#include "uid_map.h"
#undef IS_MAP_IMPLEMENTATION

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>

#include "debug.h"

#define NULL_MAP uidMapNull

#define MAP_NEW_NODE newUIDNode
#define MAP_GET_NODE getUIDNode

#define MAP_LEAF_PRED isUIDLeaf
#define MAP_RED_PRED IsRedUIDNode
#define MAP_LEFT_PRED isLeftUIDNode

#define MAP_BALANCE_INSERTION balanceUIDInsertion

#define MAP_NEAR_NODE orderedNearUIDNode
#define MAP_SIBLING_FUNC getUIDSibling
#define MAP_ROTATE_FUNC rotateUIDNode

#define LEFT true
#define RIGHT false

#define MAP_ROTATE_RIGHT(map) roateUIDNode(Map, RIGHT)
#define MAP_ROTATE_LEFT(map) rotateUIDNode(map, LEFT)

#define MAP_NEXT_NODE(map) orderedNearUIDNode(map, RIGHT)
#define MAP_PREV_NODE(map) orderedNearUIDNode(map, LEFT);

typedef bool direction;

struct MAP_STRUCT_NAME
{
    MAP_KEY key;
    
    MAP_TYPE child[2];
    
    MAP_TYPE parent;

    enum { NODE_BLACK = false, NODE_RED = true } colour;

    MAP_VALUE value;
};

static MAP_TYPE NULL_MAP = (MAP_TYPE) 0;

/*
 * Allocates a new node with the given values, the colour red
 * and two leaf children.
 */ 
static MAP_TYPE MAP_NEW_NODE(MAP_KEY key, MAP_VALUE value, MAP_TYPE parent);

/*
 * Either finds the node with the given key or the parent
 * of the node where it would be inserted.
 *
 * This function was written to maximise code reuse between
 * all operations requiering lookup.
 *
 * Note: map is required to be non-null, so any function
 * using MAP_GET_NODE has to implement a special case
 * for MAP_NULL themselves.
 */ 
static MAP_TYPE MAP_GET_NODE(MAP_KEY key, MAP_TYPE map);

/*
 * Predicate chceks if a given node is a leaf. 
 */ 
static bool MAP_LEAF_PRED(MAP_TYPE map);

/*
 * True only if the given node 'map is red.
 *
 * Note: Leaf nodes are always black.
 */ 
static bool MAP_RED_PRED(MAP_TYPE map);

/*
 * Checks if 'map' is the left child
 * of its parent.
 *
 * Note:
 * Map _can't_ be the root node, or
 * this function will fail.
 */ 
static bool MAP_LEFT_PRED(MAP_TYPE map);

/*
 * Makes sure the invariants are kept after a new node
 * is inserted.
 *
 * Since rotation may shift the root of the tree, we
 * return the root of the rebalanced tree which, unless
 * we rotate the root out of place, will be 'root'.
 *
 * Note: Maybe we should allow 'root' == 0 which would make
 * the function search for the root if needed?
 */ 
static MAP_TYPE MAP_BALANCE_INSERTION(MAP_TYPE map, MAP_TYPE root);

/*
 * Get's the next/previous in-order node of the tree, which means
 * the leftmost node of the right subtree / the rightmost node of
 * the left subtree.
 *
 * Note: The function is "curried" as MAP_NODE_NEXT/PREV.
 */
static MAP_TYPE MAP_NEAR_NODE(MAP_TYPE map, direction dir);

/*
 * Get's the sibling, that is the other child of the parent, of 'map'. 
 */ 
static MAP_TYPE MAP_SIBLING_FUNC(MAP_TYPE map);

/*
 * Does an binary-tree structure preserving rotation of 'map'
 * to the left/right.
 *
 * Note: This function is "curried" as MAP_ROTATE_LEFT/RIGHT
 */ 
static MAP_TYPE MAP_ROTATE_FUNC(MAP_TYPE map, direction dir);

/* ////////////////// EXPERIMENTAL SLAB ALLOCATOR /////////////////////// */

#ifdef MAP_USE_SLAB_ALLOCATOR

#ifndef INIT_SLAB_SIZE
#define INIT_SLAB_SIZE 2000000
#endif

static MAP_TYPE nodes = 0;
static MAP_TYPE slab = 0;

MAP_TYPE MAP_ALLOC_NODE()
{
    if(MAP_LEAF_PRED(nodes))
    {
	MAP_TYPE last;
	
	for(int x = 0; x < 1000; ++x)
	{
	    nodes = malloc((INIT_SLAB_SIZE/1000) * sizeof(struct MAP_STRUCT_NAME));
	    
	    for(int i = 0; i < ((INIT_SLAB_SIZE/1000)-1); ++i)
	    {
		nodes[i].parent = &nodes[i+1];
	    }

	    if(x != 0)
		last->parent = &nodes[0];
	    
	    last = &nodes[(INIT_SLAB_SIZE/1000) - 1];
	    
	    if(x == 0)
		slab = &nodes[0];
	}

	last->parent = NULL_MAP;
    }

    if(!MAP_LEAF_PRED(slab))
    {
	MAP_TYPE r = slab;
	slab = slab->parent;

	return r;
    }
    else
    {
	return malloc(sizeof(struct MAP_STRUCT_NAME));
    }
}

void MAP_FREE_NODE(MAP_TYPE node)
{
    node->parent = slab;
    slab = node;
}

#else
#define MAP_ALLOC_NODE() malloc(sizeof(struct MAP_STRUCT_NAME))
#define MAP_FREE_NODE free
#endif

/* /////////////// End of experimental slab allocator ////////////////// */

MAP_TYPE MAP_NEW_NODE(MAP_KEY key, MAP_VALUE val, MAP_TYPE parent)
{
    MAP_TYPE n = MAP_ALLOC_NODE();

    n->key = key;
    n->value = val;
    n->parent = parent;

    n->colour = NODE_RED;

    n->child[LEFT] = n->child[RIGHT] = MAP_NULL_FUNC();

    return n;
}

/*
 * I realise using this lookup function will result in
 * the key getting checked once more than it has to :p
 */ 
MAP_TYPE MAP_GET_NODE(MAP_KEY key, MAP_TYPE map)
{
    assert(!MAP_LEAF_PRED(map));
    
    int c = key < map->key;
    for(; (map->key != key) && !MAP_LEAF_PRED(map->child[c])
	; c = key < map->key)
    {
	map = map->child[c];
    }

    return map;
}

bool MAP_LEAF_PRED(MAP_TYPE map)
{
    return (map == NULL_MAP);
}

bool MAP_RED_PRED(MAP_TYPE map)
{
    return !MAP_LEAF_PRED(map) && map->colour;
}

bool MAP_LEFT_PRED(MAP_TYPE map)
{
    assert(!MAP_LEAF_PRED(map->parent));

    return (map == map->parent->child[LEFT]);
}

//The whole balancing stuff is still not implemented.
MAP_TYPE MAP_BALANCE_INSERTION(MAP_TYPE map, MAP_TYPE root)
{
    /*
     * This is only run when we recolour a node red,
     * thus we can assume map is red.
     */
    assert(MAP_RED_PRED(map));

    while(!MAP_LEAF_PRED(map->parent))
    {
	if(!MAP_RED_PRED(map->parent))
	{
	    /*
	     * We have a black parent, which means
	     * no invariants are broken by haveing a red
	     * node here.
	     *
	     * So we break out be returning the old root.
	     */
	    return root;
	}

	/*
	 * We know we have a red parent, which means there's
	 * a red-red clash. Balancing now depends on colour
	 * of the uncle.
	 */
	MAP_TYPE uncle = MAP_SIBLING_FUNC(map->parent);

	if(MAP_RED_PRED(uncle))
	{
	    /*
	     * We have a red uncle. Since the parent is red
	     * the grandparent has to be black. We can thus
	     * make a structure correcting transformation by
	     * "pushing" the black down one level from the
	     * grandparent into the parent and uncle and then
	     * "recurse" up the tree on the grandparent since
	     * colouring the grandparent red may violate
	     * the red-red property.
	     */
	    map->parent->colour = NODE_BLACK;
	    uncle->colour = NODE_BLACK;

	    map = map->parent->parent;
	    map->colour = NODE_RED;

	    continue;
	}

	/*
	 * We have a red node with a red parent and a black
	 * uncle. We now need some rotating. We want the
	 * case where our node and the uncle go in
	 * opposite directions.
	 *
	 * If they don't we first do a local rotation to
	 * get into that particular case.
	 */
	assert(MAP_RED_PRED(map) &&
	       MAP_RED_PRED(map->parent) &&
	       !MAP_RED_PRED(uncle));

	MAP_TYPE parent = map->parent;

	if(MAP_LEFT_PRED(map) != MAP_LEFT_PRED(map->parent))
	    parent = MAP_ROTATE_FUNC(map->parent, !MAP_LEFT_PRED(map));

	/*
	 * We now swap parent/grandparent colour and rotate
	 * the parent, "promoting" it to root.
	 *
	 * Since parent was originally red, the grandparent
	 * had to be black, thus after the rotation the only
	 * colouring difference will be one additional red
	 * node in the right subtree of the grandparent and
	 * one red node less in the left one.
	 */
	parent->colour = NODE_BLACK;
	parent->parent->colour = NODE_RED;

	map = MAP_ROTATE_FUNC(parent->parent, !MAP_LEFT_PRED(map));

	/* We're done, return the new root if there is one. */
	if(MAP_LEAF_PRED(map->parent))
	    return map;
	else
	    return root;
    }

    /*
     * We're all the way up to the root, make sure
     * it's black before we return the possibly new root.
     */
    map->colour = NODE_BLACK;
    return map;
}

MAP_TYPE MAP_BALANCED_DELETION(bool c, MAP_TYPE map, MAP_TYPE root)
{
    MAP_TYPE sibling = MAP_SIBLING_FUNC(map);
    
    /*
     * No rebalancing is needed when we reach the root node.
     */ 
    while(!MAP_LEAF_PRED(map->parent))
    {
	/*
	 * First we want to make sure map has a black sibling.
	 *
	 * If this is not the case we rotate n down and make the
	 * sibling its grandparent and swap colours between it
	 * and their parent (thus giving us a black-red-black
	 * path to map).
	 *
	 * Since the sibling was red, it can only have black
	 * children and thus maps new sibling (being a former
	 * child of the sibling) will be black.
	 *
	 * Also note that since the sibling is red, maps parent
	 * can't be red.
	 */

	if(MAP_RED_PRED(sibling))
	{
	    sibling->colour = NODE_BLACK;
	    map->parent->colour = NODE_RED;
	    
	    MAP_ROTATE_FUNC(map->parent, c);

	    /* Sibling may be the new root */
	    if(MAP_LEAF_PRED(sibling->parent))
		root = sibling;

	    /* The rotation invalidates the sibling */
	    sibling = MAP_SIBLING_FUNC(map);
	}

	/*
	 * We can now assume that map and its sibling are black.
	 *
	 * We're trying to correct the fact that all the paths
	 * through map have one black node less than the paths
	 * through its sibling.
	 * 
	 * If the siblings children are black we can recolour it
	 * red and thus reduce the number of black nodes in the
	 * paths through the sibling, though we still lack one
	 * black node in all the paths through the parent.
	 *
	 * If we can recolour the parent black (which we have to
	 * do since we coloured the sibling red) we add one black
	 * node to the paths through the parent and thus rebalance
	 * the tree (we're done).
	 *
	 * If the parent is already black, we colour it red (essentialy
	 * swapping colour with the sibling), we've now moved the problem
	 * to the parent node since all paths through the parent still
	 * have fewer black nodes than its sibling.
	 */

	if(!MAP_RED_PRED(sibling->child[LEFT]) || MAP_RED_PRED(sibling->child[RIGHT]))
	{
	    sibling->colour = NODE_RED;

	    if(MAP_RED_PRED(map->parent))
	    {
		map->parent->colour = NODE_BLACK;

		/*
		 * Maps parent can only have been rotated down the
		 * tree and can thus not be a "new" root.
		 */
		return root;
	    }
	    else
	    {
		map = map->parent;
		c = MAP_LEFT_PRED(map);
		
		continue;
	    }
	}

	/*
	 * We now know the sibling has at least one red child.
	 * We want the child in the opposite direction of map
	 * to be the red one.
	 *
	 * If this is not the case we know the other one must
	 * be black (since we've done the case where both are
	 * black). In this case we rotate "up" the other node
	 * and swap colours between the new and the old sibling.
	 */

	if(!MAP_RED_PRED(sibling->child[!c]))
	{
	    sibling = MAP_ROTATE_FUNC(sibling, c);
	    sibling->colour = NODE_BLACK;
	    sibling->child[!c]->colour = NODE_RED;

	    /*
	     * Since sibling is the sibling of map it can't
	     * be a new root.
	     */
	}

	/*
	 * We can now assume the opposite-direction child of
	 * sibling is red.
	 *
	 * We now rotate up the sibling to the root of the
	 * subtree with maps parent as root.
	 *
	 * To restore the black-depth property to the right
	 * subtree we recolour the new roots other child
	 * black (we just made sure it was red).
	 *
	 * We also swap colour between the old and the new,
	 * thus preserving the colour of the root of the
	 * subtree.
	 *
	 * Now we've added a black node through all paths
	 * through map, since the colour swap has either
	 * coloured its parent black or we've added the new
	 * root as a new black grandparent.
	 *
	 * Note that all paths that didn't go through map
	 * must have gone through its parent and sibling.
	 * Since all paths going through the subtree containing
	 * map still just go through these two which have just
	 * swapped order and colour, no properties are
	 * invalidated here.
	 *
	 * Lastly all other paths must have went through the
	 * parent, sibling and its right child (which we know
	 * was red). They now go through the sibling and the
	 * right child (which we've coloured black). The recolouring
	 * of the right child compensates for the siblings
	 * "blackness" and the sibling has the same colour
	 * as the parent (thus compensating for its possible
	 * "blackness").
	 */

	sibling->colour = map->parent->colour;
	map->parent->colour = NODE_BLACK;

	MAP_ROTATE_FUNC(map->parent, c);

	/*
	 * Since map->parent may have been root we need to
	 * check if sibling is the new root after rotation.
	 */
	if(MAP_LEAF_PRED(sibling->parent))
	    return sibling;
	else
	    return root;
    }

    /* Make sure the root is black */
    map->colour = NODE_BLACK;
    return map;
}

MAP_TYPE MAP_NEAR_NODE(MAP_TYPE map, bool dir)
{
    assert(!MAP_LEAF_PRED(map->child[dir]));
    
    for(map = map->child[dir];
	!MAP_LEAF_PRED(map->child[!dir]);
	map = map->child[!dir]) {}

    return map;
}

MAP_TYPE MAP_SIBLING_FUNC(MAP_TYPE map)
{
    /*
     * Note: Maybe we should require map to be non-root
     * and non-leaf, and make this an assert?
     */
    if(MAP_LEAF_PRED(map) || MAP_LEAF_PRED(map->parent))
	return MAP_NULL_FUNC();

    return map->parent->child[!MAP_LEFT_PRED(map)];
}

MAP_TYPE MAP_ROTATE_FUNC(MAP_TYPE map, bool dir)
{
    MAP_TYPE c = map->child[!dir];
    
    if(!MAP_LEAF_PRED(map->parent))
	map->parent->child[MAP_LEFT_PRED(map)] = c;
    
    c->parent = map->parent;
    map->parent = c;
    
    map->child[!dir] = c->child[dir];
    c->child[dir] = map;

    if(!MAP_LEAF_PRED(map->child[!dir]))
	map->child[!dir]->parent = map;
    
    return c;
}

MAP_VALUE MAP_GET_FUNC(MAP_KEY key, MAP_VALUE def, MAP_TYPE map)
{
    if(MAP_LEAF_PRED(map))
	return def; //Don't search an empty map :)

    MAP_TYPE n = MAP_GET_NODE(key, map);

    if(n->key == key)
	return n->value;
    else
	return def;
}

MAP_TYPE MAP_SET_FUNC(MAP_KEY key, MAP_VALUE val, MAP_TYPE map)
{
    if(MAP_LEAF_PRED(map))
    {
	map = MAP_NEW_NODE(key, val, MAP_NULL_FUNC());
	map->colour = NODE_BLACK;
    }

    MAP_TYPE n = MAP_GET_NODE(key, map);
    
    if(n->key == key)
    { n->value = val; return map; }
    
    n->child[key < n->key] = MAP_NEW_NODE(key, val, n);

    return MAP_BALANCE_INSERTION(n->child[key < n->key], map);
}

MAP_TYPE MAP_DEL_FUNC(MAP_KEY key, MAP_VALUE* val, MAP_TYPE map)
{
    if(MAP_LEAF_PRED(map))
	return map;

    MAP_TYPE n = MAP_GET_NODE(key, map);

    if(n->key != key)
	return map;

    /*
     * Balanced deletion depends on the deleted node having
     * at most one child. If this is not the case we swap
     * the deleted node with the next in-order node and
     * delete the next in-order node instead.
     */ 

    bool cdir;
    
    /*
     * This somewhat clunky if statement results in the
     * block only being run when we have two children,
     * and that cdir = direction the (potentialy) non-leaf
     * child is in.
     */
    if(!(cdir = MAP_LEAF_PRED(n->child[RIGHT]))
       && !MAP_LEAF_PRED(n->child[LEFT]))
    {
	cdir = RIGHT;
	
	MAP_TYPE n_ = MAP_NEXT_NODE(n);
	
	n->key = n_->key;
	n->value = n_->value;

	n = n_;
    }

    MAP_TYPE child;
    
    if(MAP_LEAF_PRED(n->child[RIGHT]) && MAP_LEAF_PRED(n->child[LEFT]))
    {
	/*
	 * TODO: These are fugly hacks, think of something better >:=O
	 */
	if(MAP_LEAF_PRED(n->parent))
	{
	    MAP_FREE_FUNC(n);
	    return MAP_NULL_FUNC();
	}
	
	if(MAP_RED_PRED(n))
	{
	    TRACE("Dump red case");
	    
	    n->parent->child[MAP_LEFT_PRED(n)] = MAP_NULL_FUNC();
	    MAP_FREE_FUNC(n);

	    return map;
	}

	cdir = !MAP_LEFT_PRED(n);

	if(MAP_RED_PRED(n->parent))
	{
	    
	}

	TRACE("Even dumber black case\n");
	
	map = MAP_BALANCED_DELETION(cdir, n->parent, map);
	MAP_FREE_FUNC(n);
	return map;
    }
    else
    {
	/*
	 * Now that we know n has one child we perform a
	 * little pointer surgery to cut n out.
	 */ 
	child = n->child[cdir];
	
	child->parent = n->parent;
	
	if(!MAP_LEAF_PRED(n->parent))
	    n->parent->child[MAP_LEFT_PRED(n)] = child;
	else
	    map = child;
    }
    
    /*
     * Simplest case is if n is red, then we just remove it
     * and replace it with its child. Since this just means
     * we pass through one less red node, it does not invalidate
     * the Black-depth property, and since the child has to be
     * black it can't invalidate the Red-Red property.
     */ 
    
    if(MAP_RED_PRED(n))
    {
	MAP_FREE_NODE(n);

	/*
	 * We know ns child cant be the new root since n
	 * is red and thus not root.
	 */ 
	return map;
    }

    /*
     * We know the node is black, deleting it will violate the
     * black depth property of the path through its child (we
     * know there's only one child).
     *
     * If the child it red we can just recolour it black to
     * restore the black-depth property.
     */

    if(MAP_RED_PRED(child))
    {
	child->colour = NODE_BLACK;

	/*
	 * We checked earlier if we've changed root
	 */
	return map;
    }

    /*
     * We had a black-black pair, requires us to do some
     * some rebalancing on the tree.
     *
     * Not that we can now assume athat child is black.
     */
    return MAP_BALANCED_DELETION(cdir, child, map);
}

bool MAP_ELEM_PRED(MAP_KEY key, MAP_TYPE map)
{
    if(MAP_LEAF_PRED(map))
	return false;
    
    MAP_TYPE n = MAP_GET_NODE(key, map);

    return (n->key == key);
}

//Note: write a nonrecursive version?
bool MAP_EQ_PRED(MAP_TYPE mapA, MAP_TYPE mapB)
{
    return (MAP_LEAF_PRED(mapA) && MAP_LEAF_PRED(mapB)) ||
	(mapA->key == mapB->key && mapA->value == mapB->value &&
	 MAP_EQ_PRED(mapA->child[LEFT], mapB->child[LEFT]) &&
	 MAP_EQ_PRED(mapA->child[RIGHT], mapB->child[RIGHT]));
}

MAP_TYPE MAP_NULL_FUNC(void)
{
    return NULL_MAP;
}

/*
 * Note: write a non-recursive version 
 */ 
void MAP_FREE_FUNC(MAP_TYPE map)
{
//    clearSlab();
    if(!MAP_LEAF_PRED(map->child[LEFT]))
	MAP_FREE_FUNC(map->child[LEFT]);
    if(!MAP_LEAF_PRED(map->child[RIGHT]))
	MAP_FREE_FUNC(map->child[RIGHT]);

    MAP_FREE_NODE(map);
}

//TEMP DEBUG
void dumpNode(MAP_TYPE map)
{
    if(MAP_LEAF_PRED(map))
	printf("O");
    else
	printf("%.3ld = %.3d", map->key, map->value);

    if(MAP_RED_PRED(map))
	printf(".r\n");
    else
	printf(".b\n");
}

void dumpMap(MAP_TYPE map, int depth)
{
    if(!depth)
	printf("Dumping tree:\n");

    for(int i = 0; i < depth; i++)
	printf("  ");
    
    dumpNode(map);

    if(MAP_LEAF_PRED(map))
	return;
    
    dumpMap(map->child[LEFT],  depth + 1);
    dumpMap(map->child[RIGHT], depth + 1);
}

bool validateStructure(MAP_TYPE map)
{
    if(!MAP_LEAF_PRED(map->child[LEFT]))
    {
	if(map->child[LEFT]->parent != map)
	{
	    TRACE("L: Error in %ld-%ld connection, thinks its parentis %ld\n",
		  map->key, map->child[LEFT]->key,
		  map->child[LEFT]->parent->key);
	    return false;
	}

	if(!validateStructure(map->child[LEFT]))
	    return false;
    }
  
    if(!MAP_LEAF_PRED(map->child[RIGHT]))
    {
	if(map->child[RIGHT]->parent != map)
	{
	    TRACE("R: Error in %ld-%ld connection, thinks its parentis %ld\n",
		  map->key, map->child[RIGHT]->key,
		  map->child[RIGHT]->parent->key);

	    return false;
	}

	if(!validateStructure(map->child[RIGHT]))
	    return false;
    }

    return true;
}

int validateMap(MAP_TYPE map)
{
    if(MAP_LEAF_PRED(map))
	return 1;

    int c = validateMap(map->child[LEFT]);
    
    assert(c == validateMap(map->child[RIGHT]));

    if(!MAP_RED_PRED(map))
	++c;

    return c;
}
